import json
import os
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
from openai import OpenAI
import httpx
import tempfile
import shutil

# 导入RAG服务
from rag_service import get_rag_service

# 初始化Flask应用
app = Flask(__name__)
CORS(app)  # 允许跨域请求

# 初始化RAG服务
get_rag_service()  # 确保RAG服务被初始化

# 初始化OpenAI客户端 - 适配OpenAI v1.x版本
# 使用自定义HTTP客户端避免proxies参数问题
custom_http_client = httpx.Client()

client = OpenAI(
    api_key="sk-457cedeb65a14e0792fdbfefd5d3d5e8",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    http_client=custom_http_client
)

@app.route('/api/chat', methods=['POST'])
def chat():
    try:
        # 获取前端发送的数据
        data = request.json
        messages = data.get('messages', [])
        stream = data.get('stream', True)  # 默认使用流式传输以提高响应速度
        
        if not messages:
            return jsonify({'error': 'No messages provided'}), 400
        
        # 获取用户最新的问题
        user_question = messages[-1]['content'] if messages and messages[-1]['role'] == 'user' else ''
        
        # 使用RAG服务检索相关文档
        rag_service = get_rag_service()
        retrieved_docs = rag_service.retrieve_documents(user_question)
        
        # 准备RAG上下文
        context = ""
        if retrieved_docs:
            # 优化上下文内容，只保留最相关的部分
            context += "根据以下相关文档内容回答问题：\n"
            # 限制上下文长度，避免过多信息影响模型响应速度
            max_context_chars = 1500
            current_length = 0
            
            for i, doc in enumerate(retrieved_docs, 1):
                doc_content = f"文档{i}: {doc['content']}\n"
                if current_length + len(doc_content) <= max_context_chars:
                    context += doc_content
                    current_length += len(doc_content)
                else:
                    # 截取部分内容，确保不超过长度限制
                    remaining_length = max_context_chars - current_length
                    context += doc_content[:remaining_length] + "...\n"
                    break
        
        # 构建带有RAG上下文的消息
        if context:
            # 创建一个新的系统消息来提供RAG上下文
            rag_messages = messages.copy()
            # 检查是否已有系统消息
            has_system_message = any(msg.get('role') == 'system' for msg in rag_messages)
            
            if has_system_message:
                # 如果已有系统消息，更新它
                for msg in rag_messages:
                    if msg.get('role') == 'system':
                        msg['content'] += "\n" + context
                        break
            else:
                # 如果没有系统消息，添加一个新的
                rag_messages.insert(0, {
                    'role': 'system',
                    'content': f"你是一个智能助手，请根据上下文信息回答用户问题。{context}\n如果提供的上下文中没有相关信息，你可以基于自己的知识回答。保持回答简洁明了。"
                })
        else:
            # 如果没有检索到文档，使用原始消息
            rag_messages = messages
        
        # 调用DeepSeek API，使用流式传输提高响应速度
        if stream:
            # 流式传输模式
            def generate():
                try:
                    # 调用API进行流式传输
                    completion = client.chat.completions.create(
                        model="deepseek-v3.2-exp",
                        messages=rag_messages,
                        stream=True  # 启用流式传输
                    )
                    
                    # 逐块发送响应内容
                    full_answer = ""
                    for chunk in completion:
                        if chunk.choices and chunk.choices[0].delta.content:
                            content = chunk.choices[0].delta.content
                            full_answer += content
                            # 以SSE格式发送数据
                            yield f"data: {json.dumps({'content': content, 'done': False})}\n\n"
                    
                    # 发送结束标记和元数据
                    yield f"data: {json.dumps({'content': '', 'done': True, 'context_used': len(retrieved_docs) > 0, 'retrieved_docs_count': len(retrieved_docs)})}\n\n"
                except Exception as e:
                    yield f"data: {json.dumps({'error': str(e), 'done': True})}\n\n"
            
            # 设置响应头，支持SSE
            return app.response_class(generate(), mimetype='text/event-stream')
        else:
            # 非流式传输模式（备用）
            completion = client.chat.completions.create(
                model="deepseek-v3.2-exp",
                messages=rag_messages,
                stream=False
            )
            
            # 处理响应
            if completion and completion.choices:
                answer = completion.choices[0].message.content
                return jsonify({
                    'success': True,
                    'answer': answer,
                    'context_used': len(retrieved_docs) > 0,
                    'retrieved_docs_count': len(retrieved_docs)
                })
            else:
                return jsonify({'error': 'No response from API'}), 500
            
    except Exception as e:
        if stream:
            # 流式传输模式下的错误处理
            def error_generate():
                yield f"data: {json.dumps({'error': str(e), 'done': True})}\n\n"
            return app.response_class(error_generate(), mimetype='text/event-stream')
        else:
            return jsonify({'error': str(e)}), 500

@app.route('/api/rag/upload', methods=['POST'])
def rag_upload():
    try:
        # 检查是否有文件上传
        if 'file' not in request.files:
            return jsonify({'error': '没有文件上传'}), 400
        
        file = request.files['file']
        
        # 检查文件名
        if file.filename == '':
            return jsonify({'error': '没有选择文件'}), 400
        
        # 保存文件到临时目录
        with tempfile.TemporaryDirectory() as temp_dir:
            file_path = os.path.join(temp_dir, file.filename)
            file.save(file_path)
            
            # 读取文件内容
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # 添加到RAG系统
            rag_service = get_rag_service()
            success = rag_service.add_document(content, {'filename': file.filename})
            
            if success:
                return jsonify({
                    'success': True,
                    'message': f'文件 {file.filename} 成功添加到RAG系统'
                })
            else:
                return jsonify({'error': '添加文件到RAG系统失败'}), 500
                
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/rag/status', methods=['GET'])
def rag_status():
    try:
        rag_service = get_rag_service()
        status = rag_service.get_vector_store_info()
        
        return jsonify({
            'success': True,
            'status': status
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/rag/add_sample', methods=['POST'])
def rag_add_sample():
    try:
        # 添加一些示例数据到RAG系统
        rag_service = get_rag_service()
        
        # 示例文档1：关于RAG技术的介绍
        rag_intro = """RAG (Retrieval-Augmented Generation) 是一种结合了检索和生成的技术。
它允许语言模型在生成回答之前检索相关的外部知识，从而提高回答的准确性和相关性。
RAG系统通常包括文档检索、文档理解和回答生成三个主要部分。"""
        
        # 示例文档2：关于南京的介绍
        nanjing_info = """南京，简称"宁"，古称金陵、建康，是江苏省省会、副省级市、南京都市圈核心城市。
南京是中国东部地区重要的中心城市、全国重要的科研教育基地和综合交通枢纽。
南京有着悠久的历史，是中国四大古都之一，有"六朝古都"、"十朝都会"之称。"""
        
        # 示例文档3：关于AI模型的介绍
        ai_model_info = """DeepSeek是一家专注于人工智能研究的公司，他们开发了多种先进的AI模型，
包括DeepSeek-R1和DeepSeek-R2等模型。这些模型在自然语言处理、计算机视觉等领域有着广泛的应用。
DeepSeek模型采用了最新的深度学习技术，能够处理复杂的任务并生成高质量的结果。"""
        
        # 添加示例文档
        rag_service.add_document(rag_intro, {'title': 'RAG技术介绍'})
        rag_service.add_document(nanjing_info, {'title': '南京城市介绍'})
        rag_service.add_document(ai_model_info, {'title': 'DeepSeek AI模型介绍'})
        
        return jsonify({
            'success': True,
            'message': '已成功添加3个示例文档到RAG系统'
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    # 在开发环境下运行，生产环境请使用WSGI服务器
    app.run(host='0.0.0.0', port=5000, debug=True)